"""
《邢不行-2023新版|Python股票量化投资课程》
author: 邢不行
微信: xbx9585

数据整理需要计算的因子脚本，可以在这里修改/添加别的因子计算
"""

import math
import numpy as np

def cal_tech_factor(df, extra_agg_dict):
    """
    计算量价因子
    :param df:
    :param extra_agg_dict:
    :return:
    """
    # =保留总市值因子
    extra_agg_dict['总市值'] = 'last'

    # =20日涨跌幅
    df['Ret_20'] = df['收盘价_复权'].pct_change(20)
    extra_agg_dict['Ret_20'] = 'last'
    df['Ret_30'] = df['收盘价_复权'].pct_change(30)
    extra_agg_dict['Ret_30'] = 'last'

    # =alpha95因子
    df['alpha95'] = df['成交额'].rolling(5).std()
    extra_agg_dict['alpha95'] = 'last'

     # ===计算换手率
    df['流通股本'] = df['流通市值'] / df['收盘价']
    df['换手率'] = df['成交量'] / df['流通股本']
    df['换手率_5'] = df['换手率'].rolling(5).mean()
    extra_agg_dict['换手率_5'] = 'last'
    df["换手率_sum"] = df["成交额"] / df["流通市值"]
    extra_agg_dict["换手率_sum"] = "sum"
    df['STR_20'] = df['换手率'].rolling(window=20, min_periods=1).std()
    extra_agg_dict['STR_20'] = 'last'

    df['换手率_20'] = df['换手率'].rolling(20).mean()
    extra_agg_dict['换手率_20'] = 'last'

     # =成交额的Std标准差
    df['成交额std_20'] = df['成交额'].rolling(window=20).std()
    extra_agg_dict['成交额std_20'] = 'last'

    # ==计算布林通道
    # ===计算指标
    # 计算均线
    # n = 20
    # df['median'] = df['收盘价'].rolling(n, min_periods=1).mean()
    # # 计算上轨、下轨道
    # df['std'] = df['收盘价'].rolling(n, min_periods=1).std(ddof=0)  # ddof代表标准差自由度
    # # df['z_score'] = abs(df['收盘价'] - df['median']) / df['std']
    # # df['m'] = df['z_score'].rolling(n, min_periods=1).mean()
    # df['m'] = 2
    # df['upper'] = df['median'] + df['m'] * df['std']
    # df['lower'] = df['median'] - df['m'] * df['std']

    # df['vol_m5'] = df['成交量'].rolling(window=5).mean()
    # df['vol_m65'] = df['成交量'].rolling(window=65).mean()

    # df['MA20斜率'] = df['median']/df['median'].shift(1)
    # df['MA20斜率'] = df['MA20斜率'].apply(lambda x : math.atan((x-1)*100)*180/3.1415926)

    # =保留成分股信息
    extra_agg_dict['沪深300成分股'] = 'last'
    extra_agg_dict['中证1000成分股'] = 'last'


    n_list = [3, 5, 10, 13, 21, 30]
    for n in n_list:
        # =n日涨跌幅
        df['Ret_%d' % n] = df['收盘价_复权'].pct_change(n)
        extra_agg_dict['Ret_%d' % n] = 'last'



        # =成交额的Std标准差
        df['成交额std_%d' % n] = df['成交额'].rolling(window=n).std()
        extra_agg_dict['成交额std_%d' % n] = 'last'



        # =新增量价相关系数_5
        df[f'量价相关系数_{n}'] = df['收盘价_复权'].rolling(n).corr(df['成交量']).shift(1)
        extra_agg_dict[f'量价相关系数_{n}'] = 'last'



        # ===JS
        df["JS_%d" % n] = (
                (df["收盘价_复权"] - df["收盘价_复权"].shift(n)) / (n * df["收盘价_复权"].shift(n)) * 100
        )
        extra_agg_dict["JS_%d" % n] = "last"



        # ===WR
        df["WR_%d" % n] = 100 * (
                (df["最高价_复权"].rolling(n, min_periods=1).max() - df["收盘价_复权"])
                / (
                        df["最高价_复权"].rolling(n, min_periods=1).max()
                        - df["最低价_复权"].rolling(n, min_periods=1).min()
                )
        )
        extra_agg_dict["WR_%d" % n] = "last"
    # ===资金类因子
    for _acc_ in ['散户', '大户', '中户', '机构']:
        df[_acc_ + '资金买入额'].fillna(value=0, inplace=True)
        df[_acc_ + '买入占比'] = df[_acc_ + '资金买入额'] * 10000 / df['成交额']
        df['总买入额_5'] = df['成交额'].rolling(5).sum()
        df[_acc_ + '资金买入额_5'] = df[_acc_ + '资金买入额'].rolling(5).sum()
        df[_acc_ + '买入占比_5'] = df[_acc_ + '资金买入额_5'] * 10000 / df['总买入额_5']
        extra_agg_dict[_acc_ + '买入占比'] = 'last'
        extra_agg_dict[_acc_ + '买入占比_5'] = 'last'



    df['主力做多'] = (df['大户买入占比'] + df['机构买入占比']) / df['散户买入占比']
    extra_agg_dict['主力做多'] = 'last'



    df["机构资金净流入"] = df["机构资金买入额"] - df["机构资金卖出额"]
    df["中户资金净流入"] = df["中户资金买入额"] - df["中户资金卖出额"]
    df["大户资金净流入"] = df["大户资金买入额"] - df["大户资金卖出额"]
    df["散户资金净流入"] = df["散户资金买入额"] - df["散户资金卖出额"]
    df["主力资金净流入"] = df["中户资金净流入"] + df["大户资金净流入"] + df["机构资金净流入"]
    df["除机构外资金净流入"] = df["中户资金净流入"] + df["大户资金净流入"] + df["散户资金净流入"]
    extra_agg_dict['机构资金净流入'] = 'last'
    extra_agg_dict['大户资金净流入'] = 'last'
    extra_agg_dict['中户资金净流入'] = 'last'
    extra_agg_dict['散户资金净流入'] = 'last'
    extra_agg_dict['主力资金净流入'] = 'last'
    extra_agg_dict['除机构外资金净流入'] = 'last'



    # =添加申万行业
    extra_agg_dict['新版申万一级行业名称'] = 'last'



    #   =试盘因子
    # df['上影线'] = (df['最高价'] - df['收盘价']) / df['收盘价']
    # df['试盘行为'] = 0
    # df.loc[(df['上影线'] >= 0.025) & (df['上影线'] <= 0.035), '试盘行为'] = 1
    # df['近期有试盘行为'] = 0
    # df.loc[(df['试盘行为'].rolling(30).sum() > 0), '近期有试盘行为'] = 1
    # df['刚有试盘行为'] = 0
    # df.loc[(df['试盘行为'].rolling(10).sum() > 0), '刚有试盘行为'] = 1
    # df['真试盘'] = 0
    # df['上影比实体'] = (df['最高价'] - df['收盘价']) / abs(df['收盘价'] - df['开盘价'])
    # df['倍量行为'] = 0
    # df.loc[(df['成交量'] > 1.4 * df['成交量'].rolling(6).mean()), '倍量行为'] = 1
    # df.loc[(df['上影比实体'] > 2.5) & (df['倍量行为'] > 0), '真试盘'] = 1
    # df['近期有真试盘行为'] = 0
    # df.loc[(df['真试盘'].rolling(30).sum() > 0), '近期有真试盘行为'] = 1

    #  奈绪选股
    # 计算均线_N:复权收盘价的N日均线，N可以取5，10，20。
    # df['均线_5'] = df['收盘价_复权'].rolling(window=5).mean()
    # df['均线_10'] = df['收盘价_复权'].rolling(window=10).mean()
    df['均线_20'] = df['收盘价_复权'].rolling(window=20).mean()
    df.loc[(df['收盘价_复权'] > df['均线_20']), 'zs_signal'] = 1
    extra_agg_dict['zs_signal'] = 'last'
    # df['均线_30'] = df['收盘价_复权'].rolling(window=30).mean()
    # df['均线_60'] = df['收盘价_复权'].rolling(window=60).mean()
    # df['均线_88'] = df['收盘价_复权'].rolling(window=88).mean()
    # df['均线_5趋势'] = df['均线_5'] - df['均线_5'].shift(1)
    # df['均线_10趋势'] = df['均线_10'] - df['均线_10'].shift(1)
    # df['均线_20趋势'] = df['均线_20'] - df['均线_20'].shift(1)
    # df['均线_30趋势'] = df['均线_30'] - df['均线_30'].shift(1)
    # df['均线_60趋势'] = df['均线_60'] - df['均线_60'].shift(1)
    # df['均线_88趋势'] = df['均线_88'] - df['均线_88'].shift(1)
    # df['均线趋势'] = df['均线_5趋势'] + df['均线_10趋势'] + df['均线_20趋势'] + df['均线_30趋势'] + df['均线_60'] + df['均线_88趋势']


    # # 计算bias_N：当日复权收盘价相对于均线_N的涨跌幅，bias_N=(收盘价_复权/均线_N)-1，N可以取5，10，20。
    # df['bias_5'] = (df['收盘价_复权'] - df['均线_5']) - 1
    # df['bias_10'] = (df['收盘价_复权'] - df['均线_10']) - 1
    # df['bias_20'] = (df['收盘价_复权'] - df['均线_20']) - 1
    # df['最高价_复权'] = df['最高价'] / df['收盘价'] * df['收盘价_复权']
    # df['最低价_复权'] = df['最低价'] / df['收盘价'] * df['收盘价_复权']
    # # 差离值, 白线
    # df['DIF'] = df['收盘价_复权'].ewm(alpha=2/13, adjust=False).mean() - df['收盘价_复权'].ewm(alpha=2/27, adjust=False).mean()
    # # 讯号线, 黄线
    # df['DEA'] = df['DIF'].ewm(alpha=2 / 10, adjust=False).mean()
    # df.loc[df['DIF'] / df['DEA'] > 1, '看涨'] = True
    # df['MACD'] = 2 * (df['DIF'] - df['DEA'])
    # df['MACD_Hist'] = abs(df['DEA'] - df['DIF'])
    # df['L97'] = 2 * (df['DIF'] - df['DIF'].ewm(alpha=2/10, adjust=False).mean())
    # # 计算N天最低价
    # df['N_min'] = df['最低价_复权'].rolling(9, min_periods=9).min()
    # df['N_min'].fillna(value=df['最低价_复权'].expanding().min(), inplace=True)
    # # 计算N天最高价
    # df['N_max'] = df['最高价_复权'].rolling(9, min_periods=9).min()
    # df['N_max'].fillna(value=df['最高价_复权'].expanding().max(), inplace=True)
    # # 计算RSV
    # df['RSV'] = (df['收盘价_复权'] - df['N_min']) / (df['N_max'] - df['N_min']) * 100
    # # 如果前边9天是NaN，填充为100
    # df['RSV'].fillna(value=100, inplace=True)
    # # 计算K、D、J的值
    # df['K'] = df['RSV'].ewm(com=2, adjust=False).mean()
    # df['D'] = df['K'].ewm(com=2, adjust=False).mean()
    # df['J'] = 3 * df['K'] - 2 * df['D']
    # # df = df[(df.中证500成分股 == 'Y')]
    # # 删除上市的第一个周期
    # # df.drop([0], axis=0, inplace=True)  # 删除第一行数据
    # # 计算所需因子
    # df['VWAP'] = df['成交额'] / df['成交量']  # 当日成交平均价格
    # df['流通股本'] = df['流通市值'] / df['收盘价']
    # df['换手率'] = df['成交量'] / df['流通股本']
    # df['成交额_10'] = df['成交额'].rolling(window=10).mean()
    # df['异常换手率'] = df['换手率'].rolling(window=10).mean() / df['换手率'].rolling(window=60).mean()
    # # df['量稳换手率变化率'] = df['换手率'].rolling(window=60).apply(lambda x: x[-20:].std() / x[:40].std() - 1)
    # df['成交额std_5'] = df['成交额'].rolling(window=5).std()
    # df['成交额std_10'] = df['成交额'].rolling(window=10).std()
    # df['成交额std_20'] = df['成交额'].rolling(window=20).std()
    # df['量价相关性'] = df['收盘价_复权'].rolling(10).corr(df['换手率'])  # 求两列的相关性
    # # df['Illiquidity'] = abs(df['涨跌幅_10']) / df['成交额_10'] * 100000000

    # 资金流筛选
    # 删除'中户资金买入额', '成交额'字段为空的行
    # df.dropna(subset=['中户资金买入额', '成交额', '散户资金卖出额', '总市值', '大户资金买入额', '收盘价', '开盘价', '前收盘价'], how='any', inplace=True, axis=0)

    # # 中单
    # df['中户资金净流入'] = df['中户资金买入额'] - df['中户资金卖出额']
    # df['mid_cash_ma3'] = df['中户资金净流入'].rolling(3).mean()
    # df['mid_cash_ma5'] = df['中户资金净流入'].rolling(5).mean()
    # df['mid_cash_ma10'] = df['中户资金净流入'].rolling(10).mean()
    # df['mid_cash_ma20'] = df['中户资金净流入'].rolling(20).mean()

    # # 大单
    # df['大户资金净流入'] = df['大户资金买入额'] - df['大户资金卖出额']
    # df['big_cash_ma3'] = df['大户资金净流入'].rolling(3).mean()
    # df['big_cash_ma5'] = df['大户资金净流入'].rolling(5).mean()
    # df['big_cash_ma10'] = df['大户资金净流入'].rolling(10).mean()
    # df['big_cash_ma20'] = df['大户资金净流入'].rolling(20).mean()

    # # 小单
    # df['散户资金净流入'] = df['散户资金买入额'] - df['散户资金卖出额']
    # df['small_cash_ma3'] = df['散户资金净流入'].rolling(3).mean()
    # df['small_cash_ma5'] = df['散户资金净流入'].rolling(5).mean()
    # df['small_cash_ma10'] = df['散户资金净流入'].rolling(10).mean()
    # df['small_cash_ma20'] = df['散户资金净流入'].rolling(20).mean()

    # # 机构单
    # df['机构资金净流入'] = df['机构资金买入额'] - df['机构资金卖出额']
    # df['super_cash_ma5'] = df['机构资金净流入'].rolling(5).mean()
    # df['super_cash_ma10'] = df['机构资金净流入'].rolling(10).mean()
    # df['super_cash_ma20'] = df['机构资金净流入'].rolling(20).mean()

    #计算AH溢价率
    if "港股收盘价" in df.columns.values:
        df.loc[(df['港股收盘价'].isna()) | (df['港股兑人民币'].isna()), 'AH溢价率'] = np.NAN
        df.loc[(df['港股收盘价'].isna() == False) & (df['港股兑人民币'].isna() == False), 'AH溢价率'] = df['AH溢价率'] = df['收盘价'] / (df['港股收盘价'] * df['港股兑人民币'])
    else:
        df['AH溢价率'] = np.NAN
    extra_agg_dict['AH溢价率'] = 'last'

    return df


def calc_fin_factor(df, extra_agg_dict):
    """
    计算财务因子
    :param df:              原始数据
    :param finance_df:      财务数据
    :param extra_agg_dict:  resample需要用到的
    :return:
    """

    # ====计算常规的财务指标
    # =归母净利润同比增速 相较于60个交易日前的变化
    df['R_np_atoopc@xbx_单季同比_60'] = df['R_np_atoopc@xbx_单季同比'].shift(60)
    df['归母净利润同比增速_60'] = df['R_np_atoopc@xbx_单季同比'] - df['R_np_atoopc@xbx_单季同比_60']
    extra_agg_dict['归母净利润同比增速_60'] = 'last'

    # ===计算单季度ROE
    df['ROE'] = df['R_np_atoopc@xbx_单季'] / df['B_total_equity_atoopc@xbx']
    extra_agg_dict['ROE'] = 'last'


     # =营业收入同比增速 相较于60个交易日前的变化
    df['R_revenue@xbx_单季同比_60'] = df['R_revenue@xbx_单季同比'].shift(60)
    df['营业收入同比增速_60'] = df['R_revenue@xbx_单季同比'] - df['R_revenue@xbx_单季同比_60']
    extra_agg_dict['营业收入同比增速_60'] = 'last'

     # ===计算存货周转率（高）
    df['存货周转率'] = df['R_operating_cost@xbx']/df['B_inventory@xbx']
    extra_agg_dict['存货周转率'] = 'last'

    # ===计算资本支出（低）
    df['资本支出'] = df['C_cash_paid_for_assets@xbx'].fillna(0)+df['C_depreciation_etc@xbx'].fillna(0)
    extra_agg_dict['资本支出'] = 'last'


    # ===计算股权融资（低）
    df['股权融资'] = df['B_capital_reserve@xbx'].fillna(0) + df['B_actual_received_capital@xbx'].fillna(0)
    extra_agg_dict['股权融资'] = 'last'

      # ===计算60个交易日前的ROE
    df['ROE_单季_60'] = df['ROE'].shift(60)
    extra_agg_dict['ROE_单季_60'] = 'last'

    # ===计算毛利率
    df['毛利率'] = (df['R_operating_total_revenue@xbx'] - df['R_operating_total_cost@xbx']) / df['R_operating_total_revenue@xbx']
    extra_agg_dict['毛利率'] = 'last'

     # 计算60个交易日（一个季度）之前的净资产
    df['B_total_equity_atoopc@xbx_60'] = df['B_total_equity_atoopc@xbx'].shift(60)
    extra_agg_dict['B_total_equity_atoopc@xbx_60'] = 'last'


    # df['R_basic_eps@xbx_单季同比_60'] = df['R_basic_eps@xbx_单季同比'].shift(60)
    # df['每股受益同比增速_60'] = df['R_basic_eps@xbx_单季同比'] - df['R_basic_eps@xbx_单季同比_60']
    # extra_agg_dict['每股受益同比增速_60'] = 'last'

    # df['R_revenue@xbx_单季同比_60'] = df['R_revenue@xbx_单季同比'].shift(60)
    # df['营业收入同比增速_60'] = df['R_revenue@xbx_单季同比'] - df['R_revenue@xbx_单季同比_60']
    # extra_agg_dict['营业收入同比增速_60'] = 'last'

    return df
